import stockcnn
import torch
import database
from torch.utils.data import TensorDataset, DataLoader
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import matplotlib.pyplot as plt


def getTestData():
    codes, end_dates, ranges, train, test = database.getTestData(1, 100)

    features = torch.tensor(train, dtype=torch.float32)
    features = features.view(-1, 1, 28, 28)
    labels = torch.tensor(test, dtype=torch.float32)
    labels = labels.view(-1, 15)
    return codes, end_dates, ranges, features, labels


def pred(ranges, out, t, i):
    outs = []
    ts = []
    for j in range(14):
        o = (
            out[j] * (ranges[i]["PriceMax"] - ranges[i]["PriceMin"])
            + ranges[i]["PriceMin"]
        )
        l = (
            t[j] * (ranges[i]["PriceMax"] - ranges[i]["PriceMin"])
            + ranges[i]["PriceMin"]
        )
        if j == 4 or j == 9:
            outs.append("-")
            ts.append("-")
        else:
            outs.append(f"{o:.3f}")
            ts.append(f"{l:.3f}")
    return outs, ts


def test():
    # 判断使用cuda
    device = "cpu"
    if torch.cuda.is_available() is True:  # cuda
        print("Cuda is available!")
        device = "cuda"
    # 实例化模型
    # cnnmodel = stockcnn.StockCnn()
    cnnmodel = torch.load("./cnnmodel-5000.model")
    cnnmodel.to(device)

    cnnmodel.eval()
    codes, end_dates, ranges, trains, labels = getTestData()

    for i in range(len(ranges)):
        data = trains[i].to(device)
        data = data.view(-1, 1, 28, 28)
        target = labels[i].numpy()
        out = cnnmodel(data)[0].to("cpu").detach().numpy()[0]
        outs, tests = pred(ranges, out, target, i)
        print("-----", codes[i], end_dates[i], "-----")
        print("o", outs)
        print("r", tests)
        print("")


test()
